#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : tngs_select_templates.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/17
@Description :
从对齐的基因序列中查找保守区域（左右引物区必须完全保守，内部区域可包含变异）。
输出用于扩增的模板序列。
"""

from argparse import ArgumentParser
from pathlib import Path
from Bio import SeqIO
import pandas as pd
from collections import defaultdict


# ----------------------------------------------------------------------
# 命令行参数
# ----------------------------------------------------------------------
def _parse_args():
    parser = ArgumentParser(description="从多序列比对中选择左右保守、内部可变的扩增区域")
    parser.add_argument("-i", "--in_dir", type=str, required=True)
    parser.add_argument("-o", "--out_dir", type=str, required=True)
    parser.add_argument("-c", "--con_len", type=int, default=50,
                        help="引物端保守长度")
    parser.add_argument("--min_amp", type=int, default=100,
                        help="扩增区最小长度")
    parser.add_argument("--max_amp", type=int, default=400,
                        help="扩增区最大长度")
    parser.add_argument("--min_genome", type=int, default=10,
                        help="模板最少覆盖的基因组数量")
    return parser.parse_args()


# ----------------------------------------------------------------------
# 是否保守（单列）
# ----------------------------------------------------------------------
def _if_conserved(series: pd.Series) -> str:
    """
    判断一个位点是否完全保守（只允许 A/T/C/G 且必须全一致）
    """
    uniq = set(series.tolist())
    if uniq in [{"A"}, {"T"}, {"C"}, {"G"}]:
        return "T"
    return "F"


# ----------------------------------------------------------------------
# 找连续的保守区（至少 con_len 长）
# ----------------------------------------------------------------------
def _select_regions(flags, con_len):
    regions = []
    n = len(flags)
    i = 0

    while i < n:
        if flags[i] == "T":
            j = i
            while j < n and flags[j] == "T":
                j += 1
            if j - i >= con_len:
                regions.append([i, j])
            i = j
        else:
            i += 1

    return regions


# ----------------------------------------------------------------------
# 合并左右保守区作为 template 区间
# ----------------------------------------------------------------------
def _merge_templates(flags, regions, min_amp, max_amp, max_snv):
    """
    regions 为保守区列表，例如 [[10,70], [200,260], ...]
    取 left=[x1,x2] 和 right=[y1,y2]
    扩增区长度 = y1 - x2（两个保守区之间的长度）
    """
    templates = []

    for i in range(len(regions)):
        left = regions[i]   # [x1, x2]
        for j in range(i + 1, len(regions)):
            right = regions[j]  # [y1, y2]

            amp_len = right[0] - left[1]
            if min_amp <= amp_len <= max_amp and flags[left[1]:right[0]].count("F") <= max_snv:
                templates.append((left, right))

    return templates


# ----------------------------------------------------------------------
# 选择最优模板序列并输出 FASTA
# ----------------------------------------------------------------------
def _write_templates(df, templates, out_dir, gene):
    """
    df: DataFrame  (position × genome)
    templates: [(left, right), ...]
    """

    out_fp = out_dir / f"{gene}.fa"
    genomes = df.columns

    with open(out_fp, "w") as w:
        for idx, (left, right) in enumerate(templates):
            best_good = -1
            best_name = ""
            best_left = ""
            best_amp = ""
            best_right = ""

            # 扩增区 = (left[1] ~ right[0]-1)
            amp_region = df.iloc[left[1]:right[0], :]

            for g in genomes:
                col = amp_region[g]

                # 非 gap 数量
                good_bp = (col != "-").sum()

                if good_bp > best_good:
                    best_good = good_bp
                    best_name = g

                    # 左右引物区 + 扩增区序列
                    left_seq = df[g].iloc[left[0]:left[1]].tolist()
                    amp_seq = df[g].iloc[left[1]:right[0]].tolist()
                    right_seq = df[g].iloc[right[0]:right[1]].tolist()

                    # 去掉 "-" gap
                    best_left = "".join([x for x in left_seq if x != "-"])
                    best_amp = "".join([x for x in amp_seq if x != "-"])
                    best_right = "".join([x for x in right_seq if x != "-"])

            # 拼接完整模板序列
            full_seq = best_left + best_amp + best_right

            w.write(
                ">Template" + str(idx + 1)
                + " | Accession=" + str(best_name)
                + " | Gene=" + str(gene)
                + " | LeftLen=" + str(left[1] - left[0])
                + " | RightLen=" + str(right[1] - right[0])
                + " | AmpLen=" + str(right[0] - left[1])
                + "\n"
                + full_seq + "\n"
            )


# ----------------------------------------------------------------------
# 主流程
# ----------------------------------------------------------------------
def main():
    args = _parse_args()

    in_dir = Path(args.in_dir).absolute()
    out_dir = Path(args.out_dir).absolute()
    out_dir.mkdir(exist_ok=True, parents=True)

    aln_files = list(in_dir.glob("*.aln"))
    if not aln_files:
        print("[ERROR] 未找到 *.aln 文件")
        return

    for aln in aln_files:
        gene = aln.stem
        print(f"[INFO] 处理基因对齐文件: {gene}")

        dic = defaultdict(list)
        counter = defaultdict(int)

        # 读取序列
        with open(aln, "r") as h:
            for rec in SeqIO.parse(h, "fasta"):
                dic[rec.id] = [s.upper() for s in rec.seq]
                counter[rec.id] += 1

        # ========= 某基因组出现多条序列 说明该基因可能存在多拷贝现象 跳过 =========
        if any(v > 1 for v in counter.values()):
            print(f"[WARN] {gene}: 某基因组出现多条序列，跳过（不适合做引物）")
            continue

        # ========= 如果这个基因覆盖的基因组数量太少，跳过 =======
        if len(counter) < args.min_genome:
            print(f"[WARN] {gene}: 覆盖的基因组不足{args.min_genome}，跳过")
            continue

        df = pd.DataFrame(dic)

        flags = df.apply(_if_conserved, axis=1).tolist()

        regions = _select_regions(flags, args.con_len)
        if not regions:
            print(f"[WARN] {gene}: 找不到满足要求的左右引物保守区")
            continue

        templates = _merge_templates(regions, args.min_amp, args.max_amp)
        if not templates:
            print(f"[WARN] {gene}: 找不到合适的扩增模板")
            continue

        _write_templates(df, templates, out_dir, gene)

        print(f"[INFO] {gene}: 完成输出\n")


if __name__ == "__main__":
    main()
